package ifly

import (
	"context"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/gorilla/websocket"
)

func NewSpark(appId, apiKey, apiSecret string) *Spark {
	return &Spark{
		endpoint:  "wss://spark-api.xf-yun.com/v2.1/chat", // wss://aichat.xf-yun.com/v1/chat
		appId:     appId,
		apiKey:    apiKey,
		apiSecret: apiSecret,
	}
}

type Spark struct {
	endpoint  string
	appId     string
	apiKey    string
	apiSecret string
}

func (r *Spark) main(ctx context.Context) {
	d := websocket.Dialer{
		HandshakeTimeout: 5 * time.Second,
	}
	conn, resp, err := d.Dial(assembleAuthUrl1(r.endpoint, r.apiKey, r.apiSecret), nil)
	if err != nil {
		panic(readResp(resp) + err.Error())
		return
	} else if resp.StatusCode != 101 {
		panic(readResp(resp) + err.Error())
	}

	go func() {
		data := genParams1(r.appId, "你是谁，可以干什么？")
		err := conn.WriteJSON(data)
		if err != nil {
			fmt.Printf("%+v\n", err)
		}
	}()

	var answer = ""
	for {
		_, msg, err := conn.ReadMessage()
		if err != nil {
			fmt.Println("read message error:", err)
			break
		}

		var data map[string]interface{}
		err1 := json.Unmarshal(msg, &data)
		if err1 != nil {
			fmt.Println("Error parsing JSON:", err)
			return
		}
		fmt.Println(string(msg))

		payload := data["payload"].(map[string]interface{})
		choices := payload["choices"].(map[string]interface{})
		header := data["header"].(map[string]interface{})
		code := header["code"].(float64)

		if code != 0 {
			fmt.Println(data["payload"])
			return
		}
		status := choices["status"].(float64)
		fmt.Println(status)
		text := choices["text"].([]interface{})
		content := text[0].(map[string]interface{})["content"].(string)
		if status != 2 {
			answer += content
		} else {
			fmt.Println("收到最终结果")
			answer += content
			usage := payload["usage"].(map[string]interface{})
			temp := usage["text"].(map[string]interface{})
			totalTokens := temp["total_tokens"].(float64)
			fmt.Println("total_tokens:", totalTokens)
			err := conn.Close()
			if err != nil {
				panic(err)
			}
			break
		}
	}

	fmt.Println(answer)

	time.Sleep(1 * time.Second)
}

func genParams1(appid, question string) map[string]interface{} { // 根据实际情况修改返回的数据结构和字段名
	messages := []SparkMessage{
		{Role: "user", Content: question},
	}

	data := map[string]interface{}{ //
		"header": map[string]interface{}{ //
			"app_id": appid, //
		},
		"parameter": map[string]interface{}{ //
			"chat": map[string]interface{}{ //
				"domain":      "general",   //
				"temperature": 0.8,         //
				"top_k":       int64(6),    //
				"max_tokens":  int64(2048), //
				"auditing":    "default",   //
			},
		},
		"payload": map[string]interface{}{ //
			"message": map[string]interface{}{ //
				"text": messages, //
			},
		},
	}
	return data //
}
func assembleAuthUrl1(hosturl string, apiKey, apiSecret string) string {
	ul, err := url.Parse(hosturl)
	if err != nil {
		fmt.Println(err)
	}

	date := time.Now().UTC().Format(time.RFC1123)
	signString := []string{"host: " + ul.Host, "date: " + date, "GET " + ul.Path + " HTTP/1.1"}
	sgin := strings.Join(signString, "\n")
	sha := HmacWithShaTobase64("hmac-sha256", sgin, apiSecret)

	authUrl := fmt.Sprintf("hmac username=\"%s\", algorithm=\"%s\", headers=\"%s\", signature=\"%s\"", apiKey, "hmac-sha256", "host date request-line", sha)
	authorization := base64.StdEncoding.EncodeToString([]byte(authUrl))

	v := url.Values{}
	v.Add("host", ul.Host)
	v.Add("date", date)
	v.Add("authorization", authorization)

	callurl := hosturl + "?" + v.Encode()
	return callurl
}
func HmacWithShaTobase64(algorithm, data, key string) string {
	mac := hmac.New(sha256.New, []byte(key))
	mac.Write([]byte(data))
	encodeData := mac.Sum(nil)
	return base64.StdEncoding.EncodeToString(encodeData)
}
func readResp(resp *http.Response) string {
	if resp == nil {
		return ""
	}
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("code=%d,body=%s", resp.StatusCode, string(b))
}
